{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp read_write_tfrecord\n",
    "import os\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Read and Write TFRecord\n",
    "\n",
    "Write bert features to TFRecord and read bert features from TFRecord."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "import json\n",
    "import os\n",
    "from fastcore.basics import partial\n",
    "from glob import glob\n",
    "from typing import Dict, Iterator, Callable\n",
    "import tempfile\n",
    "\n",
    "from loguru import logger\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from fastcore.basics import listify\n",
    "\n",
    "from m3tl.bert_preprocessing.create_bert_features import create_multimodal_bert_features\n",
    "from m3tl.special_tokens import EVAL, TRAIN\n",
    "from m3tl.params import Params\n",
    "from m3tl.utils import get_is_pyspark\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Write TFRecords"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Serialize Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "def _float_list_feature(value):\n",
    "    \"\"\"Returns a float_list from a float / double.\"\"\"\n",
    "    return tf.train.Feature(float_list=tf.train.FloatList(value=value))\n",
    "\n",
    "\n",
    "def _float_feature(value):\n",
    "    \"\"\"Returns a float_list from a float / double.\"\"\"\n",
    "    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))\n",
    "\n",
    "\n",
    "def _int64_list_feature(value):\n",
    "    \"\"\"Returns an int64_list from a bool / enum / int / uint.\"\"\"\n",
    "    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))\n",
    "\n",
    "\n",
    "def _int64_feature(value):\n",
    "    \"\"\"Returns an int64_list from a bool / enum / int / uint.\"\"\"\n",
    "    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))\n",
    "\n",
    "\n",
    "def _bytes_feature(value):\n",
    "    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))\n",
    "\n",
    "\n",
    "def _bytes_list_feature(value):\n",
    "    return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))\n",
    "\n",
    "\n",
    "def serialize_fn(features: dict, return_feature_desc=False):\n",
    "    features_tuple = {}\n",
    "    feature_desc = {}\n",
    "    for feature_name, feature in features.items():\n",
    "        if type(feature) is list:\n",
    "            feature = np.array(feature)\n",
    "        if type(feature) is np.ndarray:\n",
    "            if issubclass(feature.dtype.type, np.integer):\n",
    "                features_tuple[feature_name] = _int64_list_feature(\n",
    "                    feature.flatten())\n",
    "                feature_desc[feature_name] = 'int64'\n",
    "            # elif issubclass(feature.dtype.type, np.float):\n",
    "            else:\n",
    "                features_tuple[feature_name] = _float_list_feature(\n",
    "                    feature.flatten())\n",
    "                feature_desc[feature_name] = 'float32'\n",
    "\n",
    "            features_tuple['{}_shape'.format(\n",
    "                feature_name)] = _int64_list_feature(feature.shape)\n",
    "            feature_desc['{}_shape_value'.format(feature_name)] = feature.shape\n",
    "\n",
    "            feature_desc['{}_shape'.format(\n",
    "                feature_name)] = 'int64'\n",
    "\n",
    "            # this seems not a good idea\n",
    "            if len(feature.shape) > 1:\n",
    "                feature_desc['{}_shape_value'.format(feature_name)] = [\n",
    "                    None] + list(feature.shape[1:])\n",
    "            else:\n",
    "                feature_desc['{}_shape_value'.format(feature_name)] = [\n",
    "                    None for _ in feature.shape]\n",
    "\n",
    "        elif np.issubdtype(type(feature), np.float):\n",
    "            features_tuple[feature_name] = _float_feature(feature)\n",
    "            features_tuple['{}_shape'.format(\n",
    "                feature_name)] = _int64_list_feature([])\n",
    "            feature_desc[feature_name] = 'float32'\n",
    "\n",
    "            feature_desc['{}_shape'.format(\n",
    "                feature_name)] = 'int64'\n",
    "            feature_desc['{}_shape_value'.format(feature_name)] = []\n",
    "        elif np.issubdtype(type(feature), np.integer):\n",
    "            features_tuple[feature_name] = _int64_feature(feature)\n",
    "            features_tuple['{}_shape'.format(\n",
    "                feature_name)] = _int64_list_feature([])\n",
    "            feature_desc[feature_name] = 'int64'\n",
    "            feature_desc['{}_shape'.format(\n",
    "                feature_name)] = 'int64'\n",
    "            feature_desc['{}_shape_value'.format(feature_name)] = []\n",
    "        else:\n",
    "            if isinstance(feature, str):\n",
    "                feature = feature.encode('utf8')\n",
    "            features_tuple[feature_name] = _bytes_feature(feature)\n",
    "            features_tuple['{}_shape'.format(\n",
    "                feature_name)] = _int64_list_feature([])\n",
    "            feature_desc[feature_name] = 'string'\n",
    "            feature_desc['{}_shape'.format(\n",
    "                feature_name)] = 'int64'\n",
    "            feature_desc['{}_shape_value'.format(feature_name)] = []\n",
    "\n",
    "    example_proto = tf.train.Example(\n",
    "        features=tf.train.Features(feature=features_tuple)).SerializeToString()\n",
    "\n",
    "    if return_feature_desc:\n",
    "        return example_proto, feature_desc\n",
    "\n",
    "    return example_proto\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-19 21:40:51.722 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-19 21:40:51.723 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-19 21:40:51.723 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-19 21:40:51.723 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-19 21:40:51.724 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_regression, problem type: regression\n",
      "2021-06-19 21:40:51.724 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_fake_vector_fit, problem type: vector_fit\n",
      "2021-06-19 21:40:51.724 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-19 21:40:51.725 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem fake_contrastive_learning, problem type: contrastive_learning\n",
      "2021-06-19 21:40:51.725 | WARNING  | m3tl.base_params:assign_problem:642 - base_dir and dir_name arguments will be deprecated in the future. Please use model_dir instead.\n",
      "2021-06-19 21:40:51.726 | WARNING  | m3tl.base_params:prepare_dir:361 - bert_config not exists. will load model from huggingface checkpoint.\n",
      "2021-06-19 21:40:57.381 | INFO     | m3tl.utils:set_phase:478 - Setting phase to train\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "from m3tl.test_base import TestBase\n",
    "test_base = TestBase()\n",
    "test_features = {\n",
    "    'int_scalar': 1,\n",
    "    'float_scalar': 2.0,\n",
    "    'int_array': [1, 2, 3],\n",
    "    'float_array': np.array([4, 5, 6], dtype='float32'),\n",
    "    'int_matrix': [[1, 2, 3], [4, 5, 6]],\n",
    "    'float_matrix': np.random.uniform(size=(32, 5, 5)),\n",
    "    'string': 'this is test'\n",
    "}\n",
    "expected_desc = {'int_scalar': 'int64', 'int_scalar_shape': 'int64', 'int_scalar_shape_value': [],\n",
    "                 'float_scalar': 'float32', 'float_scalar_shape': 'int64', 'float_scalar_shape_value': [],\n",
    "                 'int_array': 'int64', 'int_array_shape_value': [None], 'int_array_shape': 'int64',\n",
    "                 'float_array': 'float32', 'float_array_shape_value': [None], 'float_array_shape': 'int64',\n",
    "                 'int_matrix': 'int64', 'int_matrix_shape_value': [None, 3], 'int_matrix_shape': 'int64',\n",
    "                 'float_matrix': 'float32', 'float_matrix_shape_value': [None, 5, 5], 'float_matrix_shape': 'int64',\n",
    "                 'string': 'string', 'string_shape': 'int64', 'string_shape_value': []}\n",
    "ser_str, feat_desc = serialize_fn(\n",
    "    features=test_features, return_feature_desc=True)\n",
    "assert feat_desc == expected_desc\n",
    "\n",
    "example = tf.train.Example()\n",
    "example.ParseFromString(ser_str)\n",
    "assert example.features.feature['int_array'].int64_list.value == [1, 2, 3]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Make TFRecord"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def make_tfrecord_local(data_list, output_dir, serialize_fn, mode='train', example_per_file=100000, prefix='', **kwargs) -> int:\n",
    "    \"\"\"\n",
    "    make tf record and return total number of records\n",
    "    \"\"\"\n",
    "    # create output tfrecord path\n",
    "    os.makedirs(os.path.join(\n",
    "        output_dir, prefix), exist_ok=True)\n",
    "\n",
    "    def _write_fn(d_list, path, serialize_fn, mode='train'):\n",
    "        logger.debug('Writing {}'.format(path))\n",
    "        feature_desc_path = os.path.join(os.path.dirname(\n",
    "            path), '{}_feature_desc.json'.format(mode))\n",
    "\n",
    "        with tf.io.TFRecordWriter(path) as writer:\n",
    "            for features in d_list:\n",
    "                example, feature_desc = serialize_fn(\n",
    "                    features, return_feature_desc=True)\n",
    "                writer.write(example)\n",
    "                if not os.path.exists(feature_desc_path):\n",
    "                    json.dump(feature_desc, open(\n",
    "                        feature_desc_path, 'w', encoding='utf8'))\n",
    "\n",
    "    _write_part_fn = partial(_write_fn, serialize_fn=serialize_fn, mode=mode)\n",
    "\n",
    "    x = []\n",
    "    total_count = 0\n",
    "    shard_count = 0\n",
    "    for idx, example in enumerate(data_list):\n",
    "        total_count += 1\n",
    "        x.append(example)\n",
    "        if idx % example_per_file == 0 and idx:  # pragma: no cover\n",
    "            path = os.path.join(\n",
    "                output_dir, prefix, '{}_{:05d}.tfrecord'.format(mode, shard_count))\n",
    "            shard_count += 1\n",
    "            _write_part_fn(d_list=x, path=path)\n",
    "            x = []\n",
    "\n",
    "    # add remaining\n",
    "    if x:\n",
    "        path = os.path.join(\n",
    "            output_dir, prefix, '{}_{:05d}.tfrecord'.format(mode, shard_count))\n",
    "        _write_part_fn(d_list=x, path=path)\n",
    "        total_count += len(x)\n",
    "    return total_count\n",
    "\n",
    "\n",
    "def make_tfrecord_pyspark(data_list, output_dir: str, serialize_fn: Callable, mode='train', example_per_file=100000, prefix='', **kwargs) -> int:\n",
    "    \"\"\"\n",
    "    make tf record and return total number of records with pyspark\n",
    "    \"\"\"\n",
    "    from m3tl.pyspark_utils import Hdfs, repar_rdd\n",
    "    from pyspark import RDD\n",
    "\n",
    "    # write RDD to TFRecords\n",
    "    # ref: https://github.com/yahoo/TensorFlowOnSpark/blob/master/examples/mnist/mnist_data_setup.py\n",
    "    # just for type hint\n",
    "    data_list: RDD = data_list\n",
    "\n",
    "    # since single record might not contain all problem labels\n",
    "    # we create feature desc for all record and aggregate\n",
    "    # TODO: poor performance, optimize this\n",
    "    feat_desc_tfrecord_tuple_rdd = data_list.map(\n",
    "        lambda x: serialize_fn(x, return_feature_desc=True)\n",
    "    )\n",
    "    feat_desc_tfrecord_tuple_rdd = feat_desc_tfrecord_tuple_rdd.cache()\n",
    "    rdd_count = int(feat_desc_tfrecord_tuple_rdd.sample(\n",
    "        False, 0.01).count() * 100)\n",
    "    feat_desc_tfrecord_tuple_rdd = repar_rdd(\n",
    "        rdd=feat_desc_tfrecord_tuple_rdd,\n",
    "        rdd_count=rdd_count,\n",
    "        example_per_par=example_per_file\n",
    "    )\n",
    "    feature_desc_pair_rdd = feat_desc_tfrecord_tuple_rdd.map(\n",
    "        lambda x: (0, x[1]))\n",
    "    tfrecord_rdd = feat_desc_tfrecord_tuple_rdd.map(\n",
    "        lambda x: (bytearray(x[0]), None))\n",
    "\n",
    "    tfrecord_rdd.saveAsNewAPIHadoopFile(\n",
    "        path=output_dir,\n",
    "        outputFormatClass=\"org.tensorflow.hadoop.io.TFRecordFileOutputFormat\",\n",
    "        keyClass=\"org.apache.hadoop.io.BytesWritable\",\n",
    "        valueClass=\"org.apache.hadoop.io.NullWritable\"\n",
    "    )\n",
    "\n",
    "    # create feature desc\n",
    "    def _update_dict(ld: dict, rd: dict) -> dict:\n",
    "        ld.update(rd)\n",
    "        return ld\n",
    "    feature_desc = feature_desc_pair_rdd.reduceByKeyLocally(_update_dict)[0]\n",
    "\n",
    "    local_feature_desc_path = '{}_feature_desc.json'.format(mode)\n",
    "    json.dump(feature_desc, open(local_feature_desc_path, 'w'), indent=4)\n",
    "    hdfs_client = Hdfs()\n",
    "    hdfs_client.copyFromLocalFile(\n",
    "        local_feature_desc_path,\n",
    "        os.path.join(output_dir, local_feature_desc_path))\n",
    "    return rdd_count\n",
    "\n",
    "\n",
    "def make_tfrecord(data_list, output_dir, serialize_fn, mode='train', example_per_file=100000, prefix='', **kwargs):\n",
    "    if get_is_pyspark():\n",
    "        output_dir = kwargs['pyspark_dir']\n",
    "        output_dir = os.path.join(output_dir, mode)\n",
    "        return make_tfrecord_pyspark(data_list=data_list,\n",
    "                                     output_dir=output_dir,\n",
    "                                     serialize_fn=serialize_fn,\n",
    "                                     mode=mode,\n",
    "                                     example_per_file=example_per_file,\n",
    "                                     prefix=prefix, **kwargs)\n",
    "\n",
    "    return make_tfrecord_local(data_list=data_list,\n",
    "                               output_dir=output_dir,\n",
    "                               serialize_fn=serialize_fn,\n",
    "                               mode=mode,\n",
    "                               example_per_file=example_per_file,\n",
    "                               prefix=prefix, **kwargs)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Local write tfrecord test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-19 21:40:57.758 | DEBUG    | __main__:_write_fn:11 - Writing /tmp/tmpa5zj7lsf/train_00000.tfrecord\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "make_tfrecord(\n",
    "    [test_features], output_dir=test_base.tmpfiledir, serialize_fn=serialize_fn)\n",
    "assert os.path.exists(os.path.join(\n",
    "    test_base.tmpfiledir, 'train_feature_desc.json'))\n",
    "assert os.path.exists(os.path.join(\n",
    "    test_base.tmpfiledir, 'train_00000.tfrecord'))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Pyspark write tfrecord test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# init spark\n",
    "from pyspark import SparkContext, SparkConf\n",
    "jar_path = '/data/m3tl/tmp/tensorflow-hadoop-1.10.0.jar'\n",
    "\n",
    "\n",
    "conf = SparkConf().set('spark.jars', jar_path)\n",
    "\n",
    "sc = SparkContext(conf=conf)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from m3tl.utils import set_is_pyspark\n",
    "import tempfile\n",
    "set_is_pyspark(True)\n",
    "test_features_rdd = sc.parallelize([test_features]).coalesce(1)\n",
    "pyspark_dir = tempfile.mkdtemp()\n",
    "make_tfrecord(\n",
    "    test_features_rdd, output_dir=test_base.tmpfiledir, serialize_fn=serialize_fn, pyspark_dir=pyspark_dir)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check data\n",
    "from m3tl.read_write_tfrecord import make_feature_desc\n",
    "json_path = os.path.join(test_base.tmpfiledir, 'train_feature_desc.json')\n",
    "feature_desc = make_feature_desc(json.load(open(json_path, 'r')))\n",
    "\n",
    "tfrecord_path = os.path.join(pyspark_dir, 'train', 'part-r-00000')\n",
    "\n",
    "tfr_dataset = tf.data.TFRecordDataset(tfrecord_path)\n",
    "\n",
    "\n",
    "def _parse_fn(x):\n",
    "    return tf.io.parse_single_example(x, feature_desc)\n",
    "\n",
    "\n",
    "tfr_dataset = tfr_dataset.map(_parse_fn)\n",
    "\n",
    "for i in tfr_dataset.take(1):\n",
    "    assert np.all(tf.sparse.to_dense(i['float_array']).numpy(\n",
    "    ) == np.array([4., 5., 6.], dtype='float32'))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Chain problems and write API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "def chain_processed_data(problem_preproc_gen_dict: Dict[str, Iterator]) -> Iterator:\n",
    "    # problem chunk size is 1, return generator directly\n",
    "    if len(problem_preproc_gen_dict) == 1:\n",
    "        return next(iter(problem_preproc_gen_dict.values()))\n",
    "\n",
    "    if get_is_pyspark():\n",
    "        from pyspark import RDD\n",
    "        from m3tl.pyspark_utils import join_dict_of_rdd\n",
    "\n",
    "        rdd = join_dict_of_rdd(rdd_dict=problem_preproc_gen_dict)\n",
    "        return rdd\n",
    "\n",
    "    logger.warning('Chaining problems with & may consume a lot of memory if'\n",
    "                   ' data is not pyspark RDD.')\n",
    "    data_dict = {}\n",
    "    column_list = []\n",
    "    for pro in problem_preproc_gen_dict:\n",
    "        data_dict[pro] = listify(problem_preproc_gen_dict[pro])\n",
    "        try:\n",
    "            column_list.append(list(data_dict[pro][0].keys()))\n",
    "        except IndexError:\n",
    "            raise IndexError(\"Problem {} has no data\".format(pro))\n",
    "\n",
    "    # get intersection and use as ensure features are the same\n",
    "    join_key = list(set(column_list[0]).intersection(*column_list[1:]))\n",
    "\n",
    "    flat_data_list = []\n",
    "    first_problem = next(iter(problem_preproc_gen_dict.keys()))\n",
    "    while data_dict[first_problem]:\n",
    "        d = {}\n",
    "        for pro in data_dict:\n",
    "            if not d:\n",
    "                d = data_dict[pro].pop(0)\n",
    "            else:\n",
    "                for k in join_key:\n",
    "                    assert d[k] == data_dict[pro][0][k], 'At iteration {}, feature {} not align. Expected {}, got: {}'.format(\n",
    "                        len(flat_data_list), k, d[k], data_dict[pro][0][k]\n",
    "                    )\n",
    "                d.update(data_dict[pro].pop(0))\n",
    "        flat_data_list.append(d)\n",
    "    return flat_data_list\n",
    "\n",
    "\n",
    "def write_tfrecord(params: Params, replace=False):\n",
    "    \"\"\"Write TFRecord for every problem chunk\n",
    "\n",
    "    Output location: params.tmp_file_dir\n",
    "\n",
    "    Arguments:\n",
    "        params {params} -- params\n",
    "\n",
    "    Keyword Arguments:\n",
    "        replace {bool} -- Whether to replace existing tfrecord (default: {False})\n",
    "    \"\"\"\n",
    "\n",
    "    read_data_fn_dict = params.read_data_fn\n",
    "    path_list = []\n",
    "    for problem_list in params.problem_chunk:\n",
    "        problem_str = '_'.join(sorted(problem_list))\n",
    "        file_dir = os.path.join(params.tmp_file_dir, problem_str)\n",
    "        if params.pyspark_output_path is not None:\n",
    "            pyspark_dir = os.path.join(params.pyspark_output_path, problem_str)\n",
    "        else:\n",
    "            pyspark_dir = None\n",
    "        if not os.path.exists(file_dir) or replace:\n",
    "            for mode in [TRAIN, EVAL]:\n",
    "\n",
    "                problem_preproc_gen_dict = {}\n",
    "                for p in problem_list:\n",
    "                    problem_preproc_gen_dict[p] = read_data_fn_dict[p](\n",
    "                        params=params, mode=mode)\n",
    "\n",
    "                chained_data = chain_processed_data(problem_preproc_gen_dict)\n",
    "\n",
    "                total_count = make_tfrecord(data_list=chained_data, output_dir=file_dir,\n",
    "                                            mode=mode, serialize_fn=serialize_fn, pyspark_dir=pyspark_dir,\n",
    "                                            example_per_file=params.example_per_file)\n",
    "                if mode == TRAIN:\n",
    "                    params.set_problem_info(\n",
    "                        problem=problem_str, info_name='data_num', info=total_count)\n",
    "                \n",
    "                if get_is_pyspark():\n",
    "                    from m3tl.pyspark_utils import Hdfs, get_text_file_from_executor\n",
    "                    # upload problem_info if pyspark\n",
    "                    local_problem_info_path = params.get_problem_info_path(problem_str)\n",
    "                    tempfile_name = tempfile.NamedTemporaryFile().name\n",
    "                    get_text_file_from_executor(local_problem_info_path, tempfile_name)\n",
    "                    params.merge_problem_info_file(tempfile_name)                    \n",
    "                    Hdfs().copyFromLocalFile(local_problem_info_path, pyspark_dir)\n",
    "                \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-19 21:41:03.855 | WARNING  | m3tl.base_params:assign_problem:642 - base_dir and dir_name arguments will be deprecated in the future. Please use model_dir instead.\n",
      "2021-06-19 21:41:03.856 | WARNING  | m3tl.base_params:prepare_dir:361 - bert_config not exists. will load model from huggingface checkpoint.\n",
      "2021-06-19 21:41:03.935 | WARNING  | __main__:chain_processed_data:15 - Chaining problems with & may consume a lot of memory if data is not pyspark RDD.\n",
      "2021-06-19 21:41:03.942 | DEBUG    | __main__:_write_fn:11 - Writing /tmp/tmpa5zj7lsf/weibo_fake_cls_weibo_fake_ner/train_00000.tfrecord\n",
      "2021-06-19 21:41:03.973 | WARNING  | __main__:chain_processed_data:15 - Chaining problems with & may consume a lot of memory if data is not pyspark RDD.\n",
      "2021-06-19 21:41:03.979 | DEBUG    | __main__:_write_fn:11 - Writing /tmp/tmpa5zj7lsf/weibo_fake_cls_weibo_fake_ner/eval_00000.tfrecord\n",
      "2021-06-19 21:41:04.005 | DEBUG    | __main__:_write_fn:11 - Writing /tmp/tmpa5zj7lsf/weibo_fake_multi_cls/train_00000.tfrecord\n",
      "2021-06-19 21:41:04.030 | DEBUG    | __main__:_write_fn:11 - Writing /tmp/tmpa5zj7lsf/weibo_fake_multi_cls/eval_00000.tfrecord\n",
      "2021-06-19 21:41:04.110 | DEBUG    | __main__:_write_fn:11 - Writing /tmp/tmpa5zj7lsf/weibo_masklm/train_00000.tfrecord\n",
      "2021-06-19 21:41:04.162 | DEBUG    | __main__:_write_fn:11 - Writing /tmp/tmpa5zj7lsf/weibo_masklm/eval_00000.tfrecord\n",
      "2021-06-19 21:41:04.229 | DEBUG    | __main__:_write_fn:11 - Writing /tmp/tmpa5zj7lsf/weibo_premask_mlm/train_00000.tfrecord\n",
      "2021-06-19 21:41:04.297 | DEBUG    | __main__:_write_fn:11 - Writing /tmp/tmpa5zj7lsf/weibo_premask_mlm/eval_00000.tfrecord\n"
     ]
    }
   ],
   "source": [
    "set_is_pyspark(False)\n",
    "test_base.params.assign_problem(\n",
    "    'weibo_fake_ner&weibo_fake_cls|weibo_fake_multi_cls|weibo_masklm|weibo_premask_mlm', base_dir=test_base.tmpckptdir)\n",
    "write_tfrecord(\n",
    "    params=test_base.params, replace=True)\n",
    "assert os.path.exists(os.path.join(\n",
    "    test_base.tmpfiledir, 'weibo_fake_cls_weibo_fake_ner'))\n",
    "assert os.path.exists(os.path.join(\n",
    "    test_base.tmpfiledir, 'weibo_fake_multi_cls'))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read TFRecords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def make_feature_desc(feature_desc_dict: dict):\n",
    "    feature_desc = {}\n",
    "    for feature_name, feature_type in feature_desc_dict.items():\n",
    "        if feature_type == 'int64':\n",
    "            feature_desc[feature_name] = tf.io.VarLenFeature(tf.int64)\n",
    "        elif feature_type == 'float32':\n",
    "            feature_desc[feature_name] = tf.io.VarLenFeature(tf.float32)\n",
    "\n",
    "    return feature_desc\n",
    "\n",
    "\n",
    "def reshape_tensors_in_dataset(example, feature_desc_dict: dict):\n",
    "    \"\"\"Reshape serialized tensor back to its original shape\n",
    "\n",
    "    Arguments:\n",
    "        example {Example} -- Example\n",
    "\n",
    "    Returns:\n",
    "        Example -- Example\n",
    "    \"\"\"\n",
    "\n",
    "    for feature_key in example:\n",
    "        example[feature_key] = tf.sparse.to_dense(example[feature_key])\n",
    "\n",
    "    @tf.function\n",
    "    def _reshape_tensor(tensor: tf.Tensor, shape_tensor: tf.Tensor, shape_tensor_in_dict: tf.Tensor):\n",
    "        \"\"\"\n",
    "        avoid empty tensor reshape error\n",
    "\n",
    "        we need to fill tensor with zeros to make sure \n",
    "        that loss multiplier aligns with features correctly\n",
    "        \"\"\"\n",
    "        if tf.equal(tf.size(tensor), 0):\n",
    "            # scalar\n",
    "            if tf.equal(tf.size(shape_tensor), 0):\n",
    "                return tf.zeros(shape=shape_tensor_in_dict, dtype=tensor.dtype)\n",
    "            else:\n",
    "                return tf.zeros(shape=shape_tensor, dtype=tensor.dtype)\n",
    "\n",
    "        return tf.reshape(tensor, shape=shape_tensor)\n",
    "\n",
    "    for feature_key in example:\n",
    "        if '_shape' in feature_key:\n",
    "            continue\n",
    "\n",
    "        shape_tensor = example['{}_shape'.format(feature_key)]\n",
    "        shape_tensor_in_dict = tf.convert_to_tensor(\n",
    "            feature_desc_dict[feature_key+'_shape_value'], dtype=tf.int32)\n",
    "\n",
    "        example[feature_key] = _reshape_tensor(\n",
    "            example[feature_key], shape_tensor, shape_tensor_in_dict)\n",
    "\n",
    "    for feature_key in list(example.keys()):\n",
    "        if '_shape' in feature_key:\n",
    "            del example[feature_key]\n",
    "\n",
    "    return example\n",
    "\n",
    "\n",
    "def add_loss_multiplier(example, problem):  # pragma: no cover\n",
    "    loss_multiplier_name = '{}_loss_multiplier'.format(problem)\n",
    "    if loss_multiplier_name not in example:\n",
    "        example[loss_multiplier_name] = tf.constant(\n",
    "            value=1, shape=(), dtype=tf.int32)\n",
    "    return example\n",
    "\n",
    "\n",
    "def set_shape_for_dataset(example, feature_desc_dict):  # pragma: no cover\n",
    "    for feature_key in example:\n",
    "        example[feature_key].set_shape(\n",
    "            feature_desc_dict['{}_shape_value'.format(feature_key)])\n",
    "    return example\n",
    "\n",
    "\n",
    "def get_dummy_features(dataset_dict, feature_desc_dict):\n",
    "    \"\"\"Get dummy features.\n",
    "    Dummy features are used to make sure every feature dict\n",
    "    at every iteration has the same keys.\n",
    "\n",
    "    Example:\n",
    "        problem A: {'input_ids': [1,2,3], 'A_label_ids': 1}\n",
    "        problem B: {'input_ids': [1,2,3], 'B_label_ids': 2}\n",
    "\n",
    "    Then dummy features:\n",
    "        {'A_label_ids': 0, 'B_label_ids': 0}\n",
    "\n",
    "    At each iteration, we sample a problem, let's say we sampled A\n",
    "    Then:\n",
    "        feature dict without dummy:\n",
    "            {'input_ids': [1,2,3], 'A_label_ids': 1}\n",
    "        feature dict with dummy:\n",
    "            {'input_ids': [1,2,3], 'A_label_ids': 1, 'B_label_ids':0}\n",
    "\n",
    "    Arguments:\n",
    "        dataset_dict {dict} -- dict of datasets of all problems\n",
    "\n",
    "    Returns:\n",
    "        dummy_features -- dict of dummy tensors\n",
    "    \"\"\"\n",
    "\n",
    "    feature_keys = [list(d.element_spec.keys())\n",
    "                    for _, d in dataset_dict.items()]\n",
    "    common_features_accross_problems = set(\n",
    "        feature_keys[0]).intersection(*feature_keys[1:])\n",
    "\n",
    "    dummy_features = {}\n",
    "    for problem, problem_dataset in dataset_dict.items():\n",
    "        output_types = {k: v.dtype for k,\n",
    "                        v in problem_dataset.element_spec.items()}\n",
    "        dummy_features.update({\n",
    "            k: tf.cast(\n",
    "                tf.constant(shape=[1 if s is None else s for s in feature_desc_dict.get('{}_shape_value'.format(k), [])],\n",
    "                            value=0),\n",
    "                v)\n",
    "            for k, v in output_types.items()\n",
    "            if k not in common_features_accross_problems})\n",
    "\n",
    "    return dummy_features\n",
    "\n",
    "\n",
    "def add_dummy_features_to_dataset(example, dummy_features):  # pragma: no cover\n",
    "    \"\"\"Add dummy features to dataset\n",
    "\n",
    "    feature dict without dummy:\n",
    "        {'input_ids': [1,2,3], 'A_label_ids': 1}\n",
    "    feature dict with dummy:\n",
    "        {'input_ids': [1,2,3], 'A_label_ids': 1, 'B_label_ids':0}\n",
    "\n",
    "    Arguments:\n",
    "        example {data example} -- dataset example\n",
    "        dummy_features {dict} -- dict of dummy tensors\n",
    "    \"\"\"\n",
    "    for feature_name in dummy_features:\n",
    "        if feature_name not in example:\n",
    "            example[feature_name] = tf.identity(dummy_features[feature_name])\n",
    "    return example\n",
    "\n",
    "\n",
    "def read_tfrecord(params: Params, mode: str):\n",
    "    \"\"\"Read and parse TFRecord for every problem\n",
    "\n",
    "    The returned dataset is parsed, reshaped, to_dense tensors\n",
    "    with dummy features.\n",
    "\n",
    "    Arguments:\n",
    "        params {params} -- params\n",
    "        mode {str} -- mode, train, eval or predict\n",
    "\n",
    "    Returns:\n",
    "        dict -- dict with keys: problem name, values: dataset\n",
    "    \"\"\"\n",
    "    dataset_dict = {}\n",
    "    all_feature_desc_dict = {}\n",
    "    for problem_list in params.problem_chunk:\n",
    "        problem = '_'.join(sorted(problem_list))\n",
    "        file_dir = os.path.join(params.tmp_file_dir, problem)\n",
    "\n",
    "        # pyspark path is different\n",
    "        local_mode_feature_desc_path = os.path.join(\n",
    "            file_dir, '{}_feature_desc.json'.format(mode))\n",
    "        if not os.path.exists(local_mode_feature_desc_path):\n",
    "            tfrecord_path_list = glob(os.path.join(\n",
    "                file_dir, mode, 'part*'))\n",
    "            feature_desc_dict = json.load(\n",
    "                open(os.path.join(file_dir, mode, '{}_feature_desc.json'.format(mode))))\n",
    "        else:\n",
    "            tfrecord_path_list = glob(os.path.join(\n",
    "                file_dir, '{}_*.tfrecord'.format(mode)))\n",
    "            feature_desc_dict = json.load(\n",
    "                open(os.path.join(file_dir, '{}_feature_desc.json'.format(mode))))\n",
    "        all_feature_desc_dict.update(feature_desc_dict)\n",
    "        feature_desc = make_feature_desc(feature_desc_dict)\n",
    "        dataset = tf.data.TFRecordDataset(\n",
    "            tfrecord_path_list, num_parallel_reads=tf.data.experimental.AUTOTUNE)\n",
    "        # when using hvd, we need to shard dataset\n",
    "        if params.use_horovod:\n",
    "            import horovod.tensorflow.keras as hvd\n",
    "            dataset = dataset.shard(hvd.size(), hvd.rank())\n",
    "        dataset = dataset.map(lambda x: tf.io.parse_single_example(\n",
    "            serialized=x, features=feature_desc), num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
    "        feature_desc_dict_replace_none = {}\n",
    "        for name, desc in feature_desc_dict.items():\n",
    "            if not isinstance(desc, list):\n",
    "                feature_desc_dict_replace_none[name] = desc\n",
    "            else:\n",
    "                desc_without_none = [i if i is not None else 1 for i in desc]\n",
    "                feature_desc_dict_replace_none[name] = desc_without_none\n",
    "\n",
    "        dataset = dataset.map(\n",
    "            lambda x: reshape_tensors_in_dataset(x, feature_desc_dict_replace_none),\n",
    "            num_parallel_calls=tf.data.experimental.AUTOTUNE).map(  # pylint: disable=no-member\n",
    "            lambda x: set_shape_for_dataset(\n",
    "                x, feature_desc_dict),\n",
    "            num_parallel_calls=tf.data.experimental.AUTOTUNE  # pylint: disable=no-member\n",
    "        )\n",
    "        for p in problem_list:\n",
    "            dataset = dataset.map(lambda x: add_loss_multiplier(x, p),\n",
    "                                  num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
    "        dataset_dict[problem] = dataset\n",
    "\n",
    "    # add dummy features\n",
    "    dummy_features = get_dummy_features(dataset_dict, all_feature_desc_dict)\n",
    "    for idx, problem in enumerate(params.get_problem_chunk(as_str=True)):\n",
    "        dataset_dict[problem] = dataset_dict[problem].map(\n",
    "            lambda x: add_dummy_features_to_dataset(x, dummy_features),\n",
    "            num_parallel_calls=tf.data.experimental.AUTOTUNE\n",
    "        )\n",
    "    return dataset_dict\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Local read tfrecord test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-19 21:41:05.884 | WARNING  | m3tl.base_params:assign_problem:642 - base_dir and dir_name arguments will be deprecated in the future. Please use model_dir instead.\n",
      "2021-06-19 21:41:05.885 | WARNING  | m3tl.base_params:prepare_dir:361 - bert_config not exists. will load model from huggingface checkpoint.\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "test_base.params.assign_problem(\n",
    "    'weibo_fake_ner&weibo_fake_cls|weibo_fake_multi_cls|weibo_masklm|weibo_premask_mlm', base_dir=test_base.tmpckptdir)\n",
    "write_tfrecord(\n",
    "    params=test_base.params, replace=False)\n",
    "dataset_dict = read_tfrecord(\n",
    "    params=test_base.params, mode='train')\n",
    "dataset: tf.data.Dataset = dataset_dict['weibo_fake_cls_weibo_fake_ner']\n",
    "assert sorted(list(dataset.element_spec.keys())) == [\n",
    "    'array_input_ids',\n",
    "    'array_mask',\n",
    "    'array_segment_ids',\n",
    "    'cate_input_ids',\n",
    "    'cate_mask',\n",
    "    'cate_segment_ids',\n",
    "    'masked_lm_ids',\n",
    "    'masked_lm_positions',\n",
    "    'masked_lm_weights',\n",
    "    'text_input_ids',\n",
    "    'text_mask',\n",
    "    'text_segment_ids',\n",
    "    'weibo_fake_cls_label_ids',\n",
    "    'weibo_fake_cls_loss_multiplier',\n",
    "    'weibo_fake_multi_cls_label_ids',\n",
    "    'weibo_fake_multi_cls_loss_multiplier',\n",
    "    'weibo_fake_ner_label_ids',\n",
    "    'weibo_fake_ner_loss_multiplier',\n",
    "    'weibo_masklm_loss_multiplier',\n",
    "    'weibo_premask_mlm_loss_multiplier',\n",
    "    'weibo_premask_mlm_masked_lm_ids',\n",
    "    'weibo_premask_mlm_masked_lm_positions',\n",
    "    'weibo_premask_mlm_masked_lm_weights']\n",
    "# make sure loss multiplier is correct\n",
    "ele = next(dataset.as_numpy_iterator())\n",
    "assert ele['weibo_fake_cls_loss_multiplier'] == 1\n",
    "assert ele['weibo_fake_ner_loss_multiplier'] == 1\n",
    "assert ele['weibo_fake_multi_cls_loss_multiplier'] == 0\n",
    "\n",
    "# multimodal dataset\n",
    "dataset: tf.data.Dataset = dataset_dict['weibo_fake_multi_cls']\n",
    "_ = next(dataset.as_numpy_iterator())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pyspark read tfrecord test\n",
    "\n",
    "NOTE: Test pyspark generated tfrecord"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-19 21:41:07.551 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem pyspark_fake_seq_tag, problem type: seq_tag\n",
      "2021-06-19 21:41:07.552 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem pyspark_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-19 21:41:07.552 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem pyspark_fake_cls, problem type: cls\n",
      "2021-06-19 21:41:07.553 | WARNING  | m3tl.base_params:assign_problem:642 - base_dir and dir_name arguments will be deprecated in the future. Please use model_dir instead.\n",
      "2021-06-19 21:41:07.554 | WARNING  | m3tl.base_params:prepare_dir:361 - bert_config not exists. will load model from huggingface checkpoint.\n",
      "2021-06-19 21:41:13.188 | INFO     | m3tl.utils:set_phase:478 - Setting phase to train\n"
     ]
    }
   ],
   "source": [
    "from m3tl.test_base import PysparkTestBase\n",
    "pyspark_test_base = PysparkTestBase()\n",
    "\n",
    "problem_chunk_str = pyspark_test_base.params.get_problem_chunk(as_str=True)[0]\n",
    "\n",
    "write_tfrecord(params=pyspark_test_base.params, replace=True)\n",
    "assert os.path.exists(os.path.join(\n",
    "    pyspark_test_base.params.pyspark_output_path, problem_chunk_str, 'problem_info.txt'))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.3 64-bit ('base': conda)",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
